import os
import subprocess
import logging
import paramiko
from getpass import getpass

# Configure logging
logging.basicConfig(filename='system_upgrade.log', level=logging.INFO, 
                    format='%(asctime)s - %(levelname)s - %(message)s')

def log_message(message):
    """Log a message to both console and log file."""
    print(message)
    logging.info(message)

def run_command(command):
    """Run a shell command and handle errors."""
    try:
        result = subprocess.run(command, check=True, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        output = result.stdout.decode('utf-8')
        log_message(output)
    except subprocess.CalledProcessError as e:
        error_message = f"Error executing command: {e.cmd}\n{e.stderr.decode('utf-8')}"
        log_message(error_message)
        exit(1)

def pull_updates(repo_path, remote_name, branch_name):
    """Pull updates from the remote repository."""
    log_message(f"Pulling updates from remote branch '{branch_name}'...")
    run_command(f"sudo ostree pull --repo={repo_path} {remote_name} {branch_name}")

def perform_upgrade():
    """Perform system upgrade."""
    log_message("Upgrading the system...")
    run_command("sudo ostree admin upgrade")

def rollback_system():
    """Rollback the system to the previous deployment."""
    log_message("Rolling back the system to the previous deployment...")
    run_command("sudo ostree admin rollback")

def update_remote_summary(remote_user, remote_server, remote_repo_path, password):
    """Update the summary file on the remote repository using paramiko for SSH."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    log_message("Connecting to remote server to update summary...")
    client.connect(remote_server, username=remote_user, password=password)
    command = f"ostree summary -u --repo={remote_repo_path}"
    stdin, stdout, stderr = client.exec_command(command)
    log_message(stdout.read().decode('utf-8'))
    log_message(stderr.read().decode('utf-8'))
    client.close()

def main():
    remote_user = "root"
    remote_server = "101.201.181.98"
    remote_repo_path = "/srv/ostree/repo"
    local_repo_path = "/ostree/repo"
    remote_name = "origin"
    branch_name = "testbranch"

    # Get password from the user
    password = getpass("Enter the password for the remote server: ")

    # Update the summary file on the remote repository
    update_remote_summary(remote_user, remote_server, remote_repo_path, password)

    # Pull updates from the remote repository
    pull_updates(local_repo_path, remote_name, branch_name)

    # Perform system upgrade
    perform_upgrade()

    # Optionally, rollback the system if needed
    # rollback_system()

if __name__ == "__main__":
    main()

